import os
import shutil
from tqdm import tqdm
import zipfile
from os.path import join as osj
from os.path import exists as ose

def extract_from_zip(zip_path,dst_path):
    # 打开ZIP文件
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # 解压文件到指定目录
            zip_ref.extractall(dst_path)
    except Exception as e:
        print(zip_path)
        os.remove( zip_path )
def move_files_to_outer_folder(folder_path,dst_folder_path):
    # 获取文件夹下的所有文件夹和文件
    for root, dirs, files in os.walk(folder_path):
        for file in tqdm(files):
            # 构建文件的原始路径和目标路径
            src_path = os.path.join(root, file)
            dest_path = os.path.join(dst_folder_path, file)
            # 移动文件
            shutil.move(src_path, dest_path)
            
            
# 三大类
# date = '0917'
# categories = ['sweater','coat',]

# 六大类
date = '1110'
categories = ['hoodie','formalatt',]


for category in categories:
# for category in ['sweater','jean']:

    # src_dir=f"/mnt/nas/datasets/diction/coat{date}"
    # src_dir=f"/mnt/nas/datasets/diction/leather{date}"
    # src_dir=f"/mnt/nas/datasets/diction/sweater{date}"
    
    src_dir=f"/mnt/nas/datasets/diction/{category}{date}"
    
    assert ose( src_dir )
    assert ose( osj( src_dir , 'names.txt' ) )
    
for category in categories:
# for category in ['coat','leather']:

    # src_dir=f"/mnt/nas/datasets/diction/coat{date}"
    # src_dir=f"/mnt/nas/datasets/diction/leather{date}"
    # src_dir=f"/mnt/nas/datasets/diction/sweater{date}"
    
    src_dir=f"/mnt/nas/datasets/diction/{category}{date}"
    
    
    dst_path=f"{src_dir}_img"
    
    
    # if ose( dst_path ):shutil.rmtree( dst_path )
    os.makedirs( dst_path , exist_ok=True )
    # 如果有 zip 先全部提取
    # 获取所有ZIP文件列表
    # zip_files = [file for file in os.listdir(src_dir) if file.endswith('.zip')]
    with open( osj(src_dir , 'names.txt'),encoding='utf-8' ) as f:
        files = f.readlines()
    zip_files = [ f.strip() for f in files if f.strip().endswith('.zip')]

    for zip_file in tqdm(zip_files):
        extract_from_zip( osj(src_dir,zip_file),dst_path )
    # 检查剩余文件夹中是否还有文件，都move出来
    # move_files_to_outer_folder(src_dir,dst_path)

'''
date=1110
bash restore.sh /mnt/nas/datasets/diction/hoodie"$date"
bash restore.sh /mnt/nas/datasets/diction/formalatt"$date"

date=1028
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img
bash restore.sh /mnt/nas/datasets/diction/jacket"$date"_img

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo_diff



'''